package ntd.util.frame.ssh;

import cn.hutool.core.io.IoUtil;
import cn.hutool.core.thread.ThreadUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.ssh.ChannelType;
import cn.hutool.extra.ssh.JschUtil;
import cn.hutool.json.JSONUtil;
import cn.hutool.log.StaticLog;
import com.jcraft.jsch.*;
import lj.netty.websocket.SocketUtils;
import ntd.util.frame.ssh.bo.SSHConnection;
import ntd.util.frame.ssh.bo.WebSSHData;
import ntd.util.util.ToolDuration;

import java.io.*;
import java.nio.charset.Charset;
import java.util.*;

public class ToolWebJsch {

    private int connectionTimeout = 5;

    private Map<String, SSHConnection> sshConnectionMap = new HashMap<>();

    public void initConnection(String key) {
        if (StrUtil.isBlank(key)) {
            return;
        }
        SSHConnection sshConnection = new SSHConnection();
        sshConnection.setJSch(new JSch());
        sshConnection.setKey(key);
        sshConnectionMap.put(key, sshConnection);
    }

    public void recvHandle(String key, String buffer) {
        WebSSHData webSSHData = JSONUtil.toBean(buffer, WebSSHData.class);

        //找到刚才存储的ssh连接对象
        SSHConnection sshConnection = sshConnectionMap.get(key);
        if (ObjectUtil.isNull(sshConnection)) {
            return;
        }

        webSSHData.setKey(key);
        if ("connect".equals(webSSHData.getOperate())) {
            // webssh连接
            ToolDuration.duration(() -> connectToSSH(sshConnection, webSSHData), "ssh 连接");
        } else if ("command".equals(webSSHData.getOperate())) {
            // webssh 命令执行
            ToolDuration.duration(() -> transToSSH(sshConnection.getChannel(), webSSHData.getCommand()),
                    StrUtil.format("命令[{}]执行", webSSHData.getCommand()));
        } else {
            StaticLog.info("不支持的操作...");
            close(key);
        }
    }

    public void sendMessage(String key, byte[] buffer) throws IOException {
        SocketUtils.push(key, new String(buffer, "utf-8"));
    }

    public void sendMessage(String key, String data) throws IOException {
        SocketUtils.push(key, data);
    }

    public void close(String key) {
        SSHConnection sshConnection = sshConnectionMap.get(key);
        if (sshConnection != null) {
            //断开连接
            JschUtil.close(sshConnection.getChannel());
            //map中移除
            sshConnectionMap.remove(key);
        }
    }

    /**
     * 使用jsch连接终端
     */
    private void connectToSSH(SSHConnection sshConnection, WebSSHData webSSHData) throws JSchException, IOException {
        //获取jsch的会话
        Session session = JschUtil.getSession(webSSHData.getHost(), webSSHData.getPort(), webSSHData.getUsername(), webSSHData.getPassword());

        if (!session.isConnected()) {
            session.connect();
        }

        Channel channel = JschUtil.openChannel(session, ChannelType.SHELL);

        //设置channel
        sshConnection.setChannel(channel);

//        //读取终端返回的信息流
        InputStream inputStream = channel.getInputStream();

        ThreadUtil.execAsync(() -> {
            try {
                sendMessage(webSSHData.getKey(), "ok");
                //循环读取
                while (channel.isConnected()) {
                    String string = StrUtil.str(readStream(inputStream), Charset.defaultCharset());
                    if (StrUtil.isNotBlank(string)) {
                        StaticLog.info("jsch 返回内容[{}]", string);
                        sendMessage(webSSHData.getKey(), string);
                    }
                    ThreadUtil.sleep(500);
                }
            } catch (IOException e) {
                throw new RuntimeException(e);
            } catch (Exception e) {
                throw new RuntimeException(e);
            } finally {
                //断开连接后关闭会话
                ToolDuration.duration(() -> {
                    IoUtil.close(inputStream);
                    JschUtil.close(channel);
//                    JschUtil.close(session);
                }, "关闭会话");
            }
        });
    }

    /**
     * 将消息转发到终端
     */
    private void transToSSH(Channel channel, String command) throws IOException, JSchException {
        if (channel != null && StrUtil.isNotBlank(command)) {
            ToolDuration.duration(() -> {
                OutputStream outputStream = channel.getOutputStream();
                if (StrUtil.endWith(command, "\r\n") || StrUtil.endWith(command, "\n")) {
                    outputStream.write((command).getBytes());
                } else {
                    outputStream.write((command + "\n").getBytes());
                }
                outputStream.flush();
            }, StrUtil.format("命令[{}]推送", command));
        }
    }

    /**
     * @param inStream
     * @return 字节数组
     * @throws Exception
     * @功能 读取流
     */
    public static byte[] readStream(InputStream inStream) throws Exception {
        byte[] b = new byte[inStream.available()];
        inStream.read(b);
        return b;
    }
}
